import os
import json
import time
import glob
from datetime import datetime
from openai import OpenAI
import threading
import queue

class DeepSeekChat:
    def __init__(self, api_key, base_url="https://api.deepseek.com/v1", model="deepseek-chat"):
        """
        初始化DeepSeek对话客户端（Stream API版本）
        """
        self.client = OpenAI(api_key=api_key, base_url=base_url)
        self.model = model
        self.conversation_history = []
        
        # 配置目录
        self.reference_dir = "reference_files"
        self.output_dir = "conversations"
        self.config_dir = "config"
        self.system_prompts_dir = os.path.join(self.config_dir, "system_prompts")
        self.setup_directories()
        
        # Stream相关变量
        self.response_queue = queue.Queue()
        self.stream_thread = None
        self.stop_stream = False
        
        # 当前选中的system prompt
        self.current_system_prompt = None
        self.current_system_prompt_name = None
        
        # 对话历史映射 (system_prompt_name -> conversation_history)
        self.conversation_histories = {}
        
        # 对话文件
        self.conversation_file = None
        self.create_conversation_file()
        
    def setup_directories(self):
        """创建必要的目录结构"""
        os.makedirs(self.reference_dir, exist_ok=True)
        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs("api_logs", exist_ok=True)
        os.makedirs(self.config_dir, exist_ok=True)
        os.makedirs(self.system_prompts_dir, exist_ok=True)
        
    def create_conversation_file(self):
        """创建对话文件"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{self.output_dir}/conversation_{timestamp}.txt"
        self.conversation_file = filename
        
        # 写入文件头
        with open(filename, "w", encoding="utf-8") as f:
            f.write("DeepSeek对话记录\n")
            f.write(f"开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("=" * 60 + "\n\n")
        
        print(f"对话记录将保存到: {filename}")
    
    def append_to_conversation_file(self, role, content):
        """追加内容到对话文件"""
        if not self.conversation_file:
            self.create_conversation_file()
            
        with open(self.conversation_file, "a", encoding="utf-8") as f:
            role_display = "用户" if role == "user" else "助手"
            f.write(f"{role_display}: {content}\n\n")
    
    def get_reference_files(self):
        """获取参考资料目录中的所有txt文件"""
        return glob.glob(os.path.join(self.reference_dir, "*.*"))
    
    def read_file_content(self, filepath):
        """读取文件内容"""
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                return f.read().strip()
        except Exception as e:
            print(f"读取文件 {filepath} 失败: {str(e)}")
            return None
    
    def build_context_prompt(self, reference_contents):
        """
        构建包含参考资料的提示词
        """
        if not reference_contents:
            return None
            
        prompt = "以下是参考资料内容，请基于这些资料回答用户的问题：\n\n"
        
        for i, ref_content in enumerate(reference_contents, 1):
            prompt += f"【参考资料 {i}】\n{ref_content}\n\n"
        
        prompt += "请根据以上参考资料回答用户的问题。如果问题与参考资料内容无关，请说明并礼貌地拒绝回答。"
        
        return prompt
    
    def load_references(self):
        """加载所有参考资料"""
        reference_files = self.get_reference_files()
        reference_contents = []
        
        for ref_file in reference_files:
            ref_content = self.read_file_content(ref_file)
            if ref_content:
                reference_contents.append(ref_content)
                
        return reference_contents, reference_files
    
    def list_system_prompts(self):
        """列出所有可用的系统提示词"""
        files = glob.glob(os.path.join(self.system_prompts_dir, "*.txt"))
        return [os.path.splitext(os.path.basename(f))[0] for f in files]
    
    def load_system_prompt(self, name):
        """加载指定名称的系统提示词"""
        system_file = os.path.join(self.system_prompts_dir, f"{name}.txt")
        if os.path.exists(system_file):
            try:
                with open(system_file, "r", encoding="utf-8") as f:
                    return f.read().strip()
            except Exception as e:
                print(f"读取系统提示词失败: {str(e)}")
                return None
        return None
    
    def save_system_prompt(self, name, prompt):
        """保存系统提示词"""
        system_file = os.path.join(self.system_prompts_dir, f"{name}.txt")
        try:
            with open(system_file, "w", encoding="utf-8") as f:
                f.write(prompt)
            print(f"系统提示词 '{name}' 已保存")
            return True
        except Exception as e:
            print(f"保存系统提示词失败: {str(e)}")
            return False
    
    def delete_system_prompt(self, name):
        """删除系统提示词"""
        system_file = os.path.join(self.system_prompts_dir, f"{name}.txt")
        if os.path.exists(system_file):
            try:
                os.remove(system_file)
                
                # 如果删除的是当前选中的提示词，清空当前选择
                if self.current_system_prompt_name == name:
                    self.current_system_prompt = None
                    self.current_system_prompt_name = None
                
                # 删除对应的对话历史
                if name in self.conversation_histories:
                    del self.conversation_histories[name]
                
                print(f"系统提示词 '{name}' 已删除")
                return True
            except Exception as e:
                print(f"删除系统提示词失败: {str(e)}")
                return False
        return False
    
    def set_current_system_prompt(self, name):
        """设置当前使用的系统提示词"""
        prompt = self.load_system_prompt(name)
        if prompt:
            self.current_system_prompt = prompt
            self.current_system_prompt_name = name
            
            # 如果该提示词没有对话历史，创建一个空列表
            if name not in self.conversation_histories:
                self.conversation_histories[name] = []
            
            # 切换到该提示词的对话历史
            self.conversation_history = self.conversation_histories[name]
            
            print(f"已切换到系统提示词: {name}")
            return True
        else:
            print(f"系统提示词 '{name}' 不存在")
            return False
    
    def get_current_system_prompt_name(self):
        """获取当前系统提示词名称"""
        return self.current_system_prompt_name
    
    def chat_stream(self, message, use_references=False, temperature=0.7, max_tokens=2000):
        """
        使用Stream API与DeepSeek模型进行对话
        """
        # 添加用户消息到历史和文件
        self.conversation_history.append({"role": "user", "content": message})
        self.append_to_conversation_file("user", message)
        
        # 清空队列和停止标志
        self.response_queue = queue.Queue()
        self.stop_stream = False
        
        # 构建消息列表
        messages = []
        
        # 添加系统提示词（如果有）
        if self.current_system_prompt:
            messages.append({"role": "system", "content": self.current_system_prompt})
        
        # 如果需要使用参考资料，添加到系统提示
        if use_references:
            reference_contents, reference_files = self.load_references()
            if reference_contents:
                system_prompt = self.build_context_prompt(reference_contents)
                if system_prompt:
                    messages.append({"role": "system", "content": system_prompt})
                    print(f"已加载 {len(reference_files)} 个参考资料")
        
        # 添加对话历史
        messages.extend(self.conversation_history)
        
        def stream_generator():
            """生成器函数用于流式处理"""
            full_response = ""
            start_time = time.time()
            token_count = 0
            
            try:
                # 创建流式响应
                stream = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    stream=True
                )
                
                # 处理流式响应
                for chunk in stream:
                    if self.stop_stream:
                        break
                        
                    if chunk.choices and chunk.choices[0].delta.content:
                        content = chunk.choices[0].delta.content
                        full_response += content
                        token_count += 1
                        self.response_queue.put(content)
                
                # 流结束，放入结束标志
                self.response_queue.put(None)
                
                # 计算响应时间
                end_time = time.time()
                response_time = end_time - start_time
                
                # 添加助手回复到历史和文件
                self.conversation_history.append({"role": "assistant", "content": full_response})
                self.append_to_conversation_file("assistant", full_response)
                
                # 记录API调用
                self.log_api_call(message, full_response, {
                    "estimated_tokens": token_count,
                    "response_time": f"{response_time:.2f}s",
                    "tokens_per_second": f"{token_count/response_time:.2f}" if response_time > 0 else "N/A",
                    "used_references": use_references,
                    "system_prompt": self.current_system_prompt_name
                })
                
            except Exception as e:
                error_msg = f"API调用出错: {str(e)}"
                self.response_queue.put(error_msg)
                self.response_queue.put(None)
        
        # 启动流式处理线程
        self.stream_thread = threading.Thread(target=stream_generator)
        self.stream_thread.daemon = True
        self.stream_thread.start()
        
        return self.get_stream_response()
    
    def get_stream_response(self):
        """
        从流中获取响应内容
        """
        while True:
            try:
                content = self.response_queue.get(timeout=30)  # 30秒超时
                if content is None:
                    break
                yield content
            except queue.Empty:
                print("\n响应超时")
                break
            except Exception as e:
                print(f"\n获取响应出错: {str(e)}")
                break
    
    def stop_current_stream(self):
        """停止当前的流式响应"""
        self.stop_stream = True
        if self.stream_thread and self.stream_thread.is_alive():
            self.stream_thread.join(timeout=1)
    
    def log_api_call(self, prompt, response, usage):
        """记录API调用日志"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = f"api_logs/api_call_{timestamp}.json"
        
        log_data = {
            "timestamp": datetime.now().isoformat(),
            "prompt": prompt[:500] + "..." if len(prompt) > 500 else prompt,
            "response_length": len(response),
            "usage": usage,
            "model": self.model
        }
        
        with open(log_file, "w", encoding="utf-8") as f:
            json.dump(log_data, f, ensure_ascii=False, indent=2)
            
    def list_reference_files(self):
        """列出参考资料文件"""
        files = self.get_reference_files()
        return [os.path.basename(f) for f in files]
    
    def print_conversation_history(self):
        """打印当前对话历史"""
        print("\n当前对话历史:")
        print("-" * 60)
        for i, msg in enumerate(self.conversation_history):
            role = "用户" if msg["role"] == "user" else "助手"
            # 只显示最近几条消息的预览
            preview = msg["content"][:100] + "..." if len(msg["content"]) > 100 else msg["content"]
            print(f"{i+1}. {role}: {preview}")
        print("-" * 60)
    
    def export_conversation(self, filename=None):
        """导出完整对话记录"""
        if not filename:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"conversation_export_{timestamp}.txt"
        
        with open(filename, "w", encoding="utf-8") as f:
            f.write("DeepSeek完整对话记录\n")
            f.write(f"导出时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("=" * 60 + "\n\n")
            
            for msg in self.conversation_history:
                role = "用户" if msg["role"] == "user" else "助手"
                f.write(f"{role}: {msg['content']}\n\n")
        
        print(f"对话已导出到: {filename}")
        return filename
    
    def clear_history(self):
        """清空当前对话历史（但保留文件记录）"""
        self.conversation_history.clear()
        print("对话历史已清空（文件记录保持不变）")
        
    def interactive_chat(self, use_references=False):
        """
        交互式聊天模式
        
        Args:
            use_references (bool): 是否使用参考资料
        """
        print(f"\n进入交互式聊天模式{'（使用参考资料）' if use_references else ''}")
        print("特殊命令: 'stop' - 停止当前响应, 'new' - 新建对话文件, 'export' - 导出对话")
        print("          'history' - 查看历史, 'clear' - 清空历史, 'back' - 返回主菜单")
        print("          'system_list' - 列出所有系统提示词")
        print("          'system_switch <名称>' - 切换到指定系统提示词")
        print("          'system_create <名称>' - 创建新系统提示词")
        print("          'system_edit <名称>' - 编辑系统提示词")
        print("          'system_delete <名称>' - 删除系统提示词")
        print("开始对话吧！")
        
        while True:
            try:
                user_input = input("\n您: ").strip()
                
                if user_input.lower() in ['back', 'exit', 'quit']:
                    self.stop_current_stream()
                    break
                elif user_input.lower() == 'stop':
                    self.stop_current_stream()
                    print("已停止当前响应")
                    continue
                elif user_input.lower() == 'new':
                    self.create_conversation_file()
                    print("已创建新的对话文件")
                    continue
                elif user_input.lower() == 'export':
                    export_file = self.export_conversation()
                    print(f"对话已导出到: {export_file}")
                    continue
                elif user_input.lower() == 'history':
                    self.print_conversation_history()
                    continue
                elif user_input.lower() == 'clear':
                    self.clear_history()
                    continue
                elif user_input.lower().startswith('system_list'):
                    prompts = self.list_system_prompts()
                    if prompts:
                        print("\n可用系统提示词:")
                        for i, name in enumerate(prompts, 1):
                            print(f"{i}. {name}")
                    else:
                        print("没有可用的系统提示词")
                    continue
                elif user_input.lower().startswith('system_switch'):
                    parts = user_input.split()
                    if len(parts) >= 2:
                        name = parts[1]
                        self.set_current_system_prompt(name)
                    else:
                        print("请指定要切换的系统提示词名称")
                    continue
                elif user_input.lower().startswith('system_create'):
                    parts = user_input.split()
                    if len(parts) >= 2:
                        name = parts[1]
                        prompt = input("请输入系统提示词内容: ").strip()
                        if prompt:
                            self.save_system_prompt(name, prompt)
                            self.set_current_system_prompt(name)
                    else:
                        print("请指定新系统提示词的名称")
                    continue
                elif user_input.lower().startswith('system_edit'):
                    parts = user_input.split()
                    if len(parts) >= 2:
                        name = parts[1]
                        current_content = self.load_system_prompt(name)
                        if current_content is not None:
                            print(f"当前内容: {current_content}")
                            new_content = input("请输入新内容: ").strip()
                            if new_content:
                                self.save_system_prompt(name, new_content)
                                if self.current_system_prompt_name == name:
                                    self.current_system_prompt = new_content
                        else:
                            print(f"系统提示词 '{name}' 不存在")
                    else:
                        print("请指定要编辑的系统提示词名称")
                    continue
                elif user_input.lower().startswith('system_delete'):
                    parts = user_input.split()
                    if len(parts) >= 2:
                        name = parts[1]
                        self.delete_system_prompt(name)
                    else:
                        print("请指定要删除的系统提示词名称")
                    continue
                elif not user_input:
                    continue
                
                # 使用Stream API进行对话
                print("\n助手: ", end="", flush=True)
                
                full_response = ""
                for chunk in self.chat_stream(user_input, use_references=use_references):
                    print(chunk, end="", flush=True)
                    full_response += chunk
                
                print()  # 换行
                
            except KeyboardInterrupt:
                self.stop_current_stream()
                print("\n\n返回主菜单")
                break
            except Exception as e:
                print(f"\n发生错误: {str(e)}")
                break